import glob
import json
import logging
import os,sys,time
from rknn.api import RKNN
import argparse

import common_utils
import repo_uploader
from config import generate_config
from logger import ULogger
from common_utils import get_modify_time,check_modify_time,modify_load_config
import shutil

def prepare_image_list(work_dir,cache_dir):
    img_exts=[".jpeg",".jpg",".png"]
    image_dir=os.path.join(work_dir,"quant_imgs")
    image_lst=os.listdir(image_dir)
    images=[]
    for image in image_lst:
        ext=os.path.splitext(image)
        if ext in img_exts:
            images.append(os.path.join(os.path.abspath(image_dir),image))

    with open(os.path.join(cache_dir,"image_list.txt"),'w') as f:
        f.writelines(images)

def check_metadata(configs):
    meta_path=os.path.join(configs.work_dir,"meta.json")
    if os.path.exists(meta_path):
        shutil.copy(meta_path,os.path.join(configs.cache_repo,"meta.json"))
    meta_path=os.path.join(configs.cache_repo,"meta.json")
    requant=False
    try:
        with open(meta_path,'r') as f:
            meta_json=json.load(f)
    except Exception as e:
        meta_json=dict()

    #check convert json's time
    requant = check_modify_time(os.path.join(configs.work_dir,"convert_config.json"),meta_json)

    #check model modify time
    if configs.source_framework == "caffe":
        proto = os.path.join(configs.work_dir, os.path.pardir, "caffe", configs.model_version + ".prototxt")
        model = proto.replace(".prototxt",".caffemodel")
        if not os.path.exists(proto) or not os.path.exists(model):
            requant =False
        else:
            requant &= (check_modify_time(proto,meta_json) and check_modify_time(model,meta_json))

        ## todo: onnx, mxnet
        # save meta.json
    with open(meta_path,'w+') as f:
        json.dump(meta_json,f)
        #check file exist
    requant &= (not os.path.exists(os.path.join(configs.work_dir,configs.model_version+".rknn")))
    requant &= (os.path.exists(os.path.join(configs.work_dir,"quant_imgs")))
    return requant

def upload_products(configs):
    ### upload
    #upload rknn model
    uploader=repo_uploader.RepoUploader()
    export_path = os.path.join(configs.cache_repo, configs.model_version + ".rknn")
    uploader.upload_object(export_path,configs.model_name,os.path.join(configs.model_version,"rknn",configs.model_version+".rknn"))
    #upload log
    uploader.upload_object(os.path.join(configs.cache_repo,"quant.log"),configs.model_name,os.path.join(configs.model_version,"rknn","quant.log"))
    #upload meta data
    uploader.upload_object(os.path.join(configs.cache_repo,"meta.json"),configs.model_name,os.path.join(configs.model_version,"rknn","meta.json"))

    uploader.upload_object(os.path.join(configs.cache_repo,os.path.pardir,os.path.pardir,"load_config.json"),configs.model_name,"load_config.json")
    #remove cache
    os.system("rm -rf {0}".format(configs.cache_repo))
    # shutil.rmtree(configs.cache_repo,ignore_errors=True)

def quant(configs):
    #create cache folder
    configs.cache_repo=os.path.join(configs.cache_repo,configs.model_name,configs.model_version,"rknn")
    print(configs.cache_repo)
    os.makedirs(os.path.join(configs.cache_repo),exist_ok=True)
    ulogger = ULogger(filename=os.path.join(configs.cache_repo, "quant.log"), level=logging.INFO)
    if not check_metadata(configs):
        ulogger.info("task skipped: {0} {1} {2}".format(configs.model_name, configs.model_version, "rknn"))
        return
    
    ulogger.info("start quant")
    rknn = RKNN(verbose=True,verbose_file=os.path.join(configs.cache_repo,"quant_backup.log"))
    rknn_logger = rknn.rknn_log
    log_level={
        "debug":10,
        "info":20,
        "warn":30,
        "error":40,
        "critical":50
    }
    if configs.log_level in log_level:
        rknn_logger.logger.level=log_level[configs.log_level]
    else:
        ulogger.warning("configured log level is invalid")
    std_vals=configs.scale
    if std_vals[0]<1:
        std_vals = [1/val for val in std_vals]

    reorder = ' '.join([str(i) for i in configs.reorder])
    rknn.config(mean_values=[configs.mean],
                std_values=[std_vals],
                reorder_channel=reorder,
                target_platform=['rv1109'])

    #prepare imgs
    prepare_image_list(configs.work_dir,configs.cache_repo)
    if configs.source_framework == "caffe":
        model=os.path.join(configs.work_dir,os.path.pardir,"caffe",configs.model_version+".prototxt")
        ret = rknn.load_caffe(model=model,blobs = model.replace(".prototxt",".caffemodel"),proto="caffe")
        if ret !=0:
            raise Exception("load caffe failed {0}".format(ret))
    elif configs.source_framework == "onnx":
        model = os.path.join(configs.work_dir, os.path.pardir, "onnx", configs.model_version + ".onnx")
        ret = rknn.load_onnx(model=model)
        if ret != 0:
            raise Exception("load onnx failed {0}".format(ret))
    elif configs.source_framework == "mxnet":
        params = glob.glob(os.path.join(configs.work_dir, os.path.pardir, "mxnet","{0}-*.params".format(configs.model_version)))
        if params is None or len(params)!=1:
            raise Exception("search mxnet model failed")
        symbol=os.path.join(configs.work_dir, os.path.pardir, "mxnet","{0}.symbol".format(configs.model_version))
        rknn.load_mxnet(symbol=symbol,params=params[0],input_size_list=[configs.shape])
    else:
        ulogger.error("unsupported framework {0}".format(configs.source_framework))
        raise NotImplementedError("unsupported framework")
    ret = rknn.build(do_quantization=True,dataset=os.path.join(configs.cache_repo,"image_list.txt"),pre_compile=True)
    if ret != 0:
        ulogger.error("build failed with ret {0}".format(ret))
        raise Exception("build failed: {0}".format(ret))

    export_path=os.path.join(configs.cache_repo,configs.model_version+".rknn")
    ret = rknn.export_rknn(export_path=export_path)
    if ret != 0:
        ulogger.error("build failed with ret {0}".format(ret))
        raise Exception("export failed: {0}".format(ret))

    if configs.load_after_convert:
        common_utils.modify_load_config(configs)

    upload_products(configs)

def run_quant(config):
    try:
        configs = generate_config(config)
    except FileNotFoundError:
        return 0
    except Exception as e:
        logging.info("quant failed: {0}".format(e))
        return -1

    try:
        quant(configs)
    except Exception as e:
        logging.error("quant failed: {0}".format(e))
        upload_products(configs)
        return -1
    return 0

if __name__=="__main__":
    argparser=argparse.ArgumentParser()
    argparser.add_argument("config",type=str)
    args = argparser.parse_args()
    try:
        configs = generate_config(args.config)
    except FileNotFoundError:
        exit(0)
    except Exception as e:
        logging.info("quant failed: {0}".format(e))
        exit(-1)

    try:
        quant(configs)
    except Exception as e:
        logging.error("quant failed: {0}".format(e))
        upload_products(configs)
        exit(-1)


